import sqlite3
import os
import sys
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

# 获取数据库路径
ROOT_DIR = Path(__file__).parent.parent
DB_PATH = ROOT_DIR / "articles.db"

def update_prompt_templates_table():
    """更新提示词模板表结构"""
    conn = sqlite3.connect(str(DB_PATH))
    cursor = conn.cursor()
    
    # 检查表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='prompt_templates'")
    if not cursor.fetchone():
        print("提示词模板表不存在，将创建新表")
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS prompt_templates (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                description TEXT,
                system_prompt TEXT,
                user_prompt TEXT NOT NULL,
                variables TEXT,  -- JSON 格式存储变量列表
                category TEXT,
                tags TEXT,  -- JSON 格式存储标签列表
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
    else:
        print("提示词模板表已存在，将更新表结构")
        
        # 创建新表
        cursor.execute("""
            CREATE TABLE IF NOT EXISTS prompt_templates_new (
                id TEXT PRIMARY KEY,
                name TEXT NOT NULL,
                description TEXT,
                system_prompt TEXT,
                user_prompt TEXT NOT NULL,
                variables TEXT,  -- JSON 格式存储变量列表
                category TEXT,
                tags TEXT,  -- JSON 格式存储标签列表
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
        """)
        
        # 获取现有列
        cursor.execute("PRAGMA table_info(prompt_templates)")
        columns = [column[1] for column in cursor.fetchall()]
        print(f"现有列: {columns}")
        
        # 构建复制数据的SQL
        common_columns = []
        for col in columns:
            if col in ['id', 'name', 'description', 'system_prompt', 'user_prompt', 
                      'variables', 'category', 'tags', 'created_at', 'updated_at']:
                common_columns.append(col)
        
        # 如果没有 user_prompt 列，使用 template 列（如果存在）
        if 'user_prompt' not in common_columns and 'template' in columns:
            # 复制数据时将 template 列的值复制到 user_prompt 列
            copy_columns = [col if col != 'template' else 'template as user_prompt' 
                           for col in common_columns if col != 'user_prompt']
            copy_columns_str = ', '.join(copy_columns)
        else:
            copy_columns_str = ', '.join(common_columns)
        
        if copy_columns_str:
            try:
                # 复制数据
                cursor.execute(f"INSERT INTO prompt_templates_new ({copy_columns_str}) SELECT {copy_columns_str} FROM prompt_templates")
                print(f"已复制数据: {cursor.rowcount} 行")
                
                # 删除旧表
                cursor.execute("DROP TABLE prompt_templates")
                print("已删除旧表")
                
                # 重命名新表
                cursor.execute("ALTER TABLE prompt_templates_new RENAME TO prompt_templates")
                print("已重命名新表")
            except sqlite3.Error as e:
                print(f"复制数据时出错: {e}")
                # 回滚事务
                conn.rollback()
                # 删除临时表
                cursor.execute("DROP TABLE IF EXISTS prompt_templates_new")
                print("已删除临时表")
                return
        else:
            print("没有可复制的列，将直接替换表")
            cursor.execute("DROP TABLE prompt_templates")
            cursor.execute("ALTER TABLE prompt_templates_new RENAME TO prompt_templates")
    
    conn.commit()
    conn.close()
    
    print(f"提示词模板表结构已更新: {DB_PATH}")

if __name__ == "__main__":
    update_prompt_templates_table() 